package com.flute.haflute.jobbox.base.storage;

import static me.prettyprint.cassandra.utils.StringUtils.string;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import me.prettyprint.cassandra.service.CassandraClient;
import me.prettyprint.cassandra.service.CassandraClientPool;
import me.prettyprint.cassandra.service.CassandraClientPoolFactory;
import me.prettyprint.cassandra.service.Keyspace;

import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnParent;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;

public class EqualitySelector {
	private String keySpaceName = "Keyspace1";
	private String cfName = "Standard1";
	private CassandraClientPool pool;
	private CassandraClient client;
	private SlicePredicate sp;
	private ColumnParent columnParent;
	private Keyspace keyspace;

	public EqualitySelector(String cassandraServer, String keySpace, String columnFamily) throws Exception {
		this(cassandraServer);
		this.keySpaceName = keySpace;
		this.cfName = columnFamily;
	}

	public EqualitySelector(String cassandraServer) throws Exception {
		pool = CassandraClientPoolFactory.INSTANCE.get();
		client = pool.borrowClient(cassandraServer);
		initPredicate();
	}
	
	public EqualitySelector(CassandraClient client) throws Exception {
		this.client = client;
		initPredicate();
	}

	private void initPredicate() throws Exception {
		keyspace = client.getKeyspace(keySpaceName);
		columnParent = new ColumnParent(cfName);		
		sp = new SlicePredicate();
		SliceRange sr = new SliceRange();
		sr.setStart(new byte[0]);
		sr.setFinish(new byte[0]);
		sp.setSlice_range(sr);
	}
	
	public Map<String, String> query(String key) throws Exception {
		Map<String, String> result = new HashMap<String, String>();
		List<Column> row = keyspace.getSlice(key, columnParent, sp);
		for(Column col : row) {
			result.put(string(col.getName()), string(col.getValue()));
		}
		return result;
	}
	
	public Map<String, Map<String, String>> query(List<String> keyList) throws Exception {
		Map<String, Map<String, String>> result = new HashMap<String, Map<String, String>>(keyList.size());
		for(String key : keyList) {
			result.put(key, query(key));
		}
		return result;		
	}
	
	public void close() {
		if (pool != null) {
			try {
				pool.releaseClient(client);
				keyspace = null;
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
		
	}
}
